import sys
import cv2
import os
import numpy as np
from datetime import datetime
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *


class FaceRecognitionSystem(QWidget):
    def __init__(self):
        super().__init__()
        self.cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.label_map = {}
        self.frame_counter = 0
        self.is_paused = False
        self.conf_threshold = 80  # 默认置信度阈值

        self.init_ui()
        self.load_dataset()
        self.timer = self.startTimer(30)

    def init_ui(self):
        """优化后的现代风格界面"""
        self.setWindowTitle("智能人脸识别考勤系统")
        self.setGeometry(100, 100, 1024, 600)
        self.setStyleSheet("""
            QWidget {
                background-color: #f0f2f5;
                font-family: "微软雅黑";
            }
            QLabel {
                color: #2c3e50;
            }
            QPushButton {
                background-color: #3498db;
                color: white;
                border-radius: 5px;
                padding: 8px 20px;
                font-size: 14px;
            }
            QPushButton:hover {
                background-color: #2980b9;
            }
            QLineEdit {
                border: 2px solid #bdc3c7;
                border-radius: 5px;
                padding: 8px;
                font-size: 14px;
            }
        """)

        main_layout = QHBoxLayout()

        # 左侧视频显示区（占70%宽度）
        video_panel = QWidget()
        video_layout = QVBoxLayout()
        self.video_label = QLabel()
        self.video_label.setFixedSize(700, 500)
        self.video_label.setStyleSheet("background-color: #2c3e50; border-radius: 10px;")
        video_layout.addWidget(self.video_label, alignment=Qt.AlignCenter)
        video_panel.setLayout(video_layout)
        main_layout.addWidget(video_panel, 7)

        # 右侧控制面板（占30%宽度）
        control_panel = QWidget()
        control_layout = QVBoxLayout()
        control_layout.setSpacing(20)
        control_layout.setContentsMargins(20, 20, 20, 20)

        # 输入框组
        input_group = QGroupBox("用户信息")
        input_layout = QFormLayout()
        self.name_edit = QLineEdit()
        self.name_edit.setPlaceholderText("输入姓名...")
        self.id_edit = QLineEdit()
        self.id_edit.setPlaceholderText("输入学号...")
        input_layout.addRow("姓名：", self.name_edit)
        input_layout.addRow("学号：", self.id_edit)
        input_group.setLayout(input_layout)
        control_layout.addWidget(input_group)

        # 控制按钮组
        btn_group = QWidget()
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(10)
        self.play_pause_btn = QPushButton("开始识别")
        self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-start"))
        self.play_pause_btn.clicked.connect(self.toggle_play_pause)
        self.capture_btn = QPushButton("采集人脸")
        self.capture_btn.setIcon(QIcon.fromTheme("camera-photo"))
        self.capture_btn.clicked.connect(self.capture_face)
        self.punch_btn = QPushButton("立即打卡")
        self.punch_btn.setIcon(QIcon.fromTheme("system-run"))
        self.punch_btn.clicked.connect(self.punch_in)
        btn_layout.addWidget(self.play_pause_btn)
        btn_layout.addWidget(self.capture_btn)
        btn_layout.addWidget(self.punch_btn)
        btn_group.setLayout(btn_layout)
        control_layout.addWidget(btn_group)

        # 状态显示区
        status_group = QGroupBox("系统状态")
        status_layout = QVBoxLayout()
        self.status_label = QLabel("就绪 | 摄像头正常")
        self.status_label.setStyleSheet("color: #27ae60; font-size: 16px;")
        self.conf_slider = QSlider(Qt.Horizontal)
        self.conf_slider.setRange(60, 95)
        self.conf_slider.setValue(self.conf_threshold)
        self.conf_slider.valueChanged.connect(self.update_threshold)
        status_layout.addWidget(self.status_label)
        status_layout.addWidget(QLabel("识别灵敏度："))
        status_layout.addWidget(self.conf_slider)
        status_group.setLayout(status_layout)
        control_layout.addWidget(status_group)

        control_panel.setLayout(control_layout)
        main_layout.addWidget(control_panel, 3)

        self.setLayout(main_layout)

    def update_frame(self):
        """优化后的视频渲染流程"""
        if self.is_paused:
            return

        ret, frame = self.cap.read()
        if not ret:
            self.status_label.setText("错误：无法访问摄像头！")
            self.status_label.setStyleSheet("color: #e74c3c;")
            return

        # 人脸检测与绘制
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        for (x, y, w, h) in faces:
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            # 实时识别（仅在识别模式时显示）
            if self.play_pause_btn.text() == "停止识别":
                face_roi = gray[y:y + h, x:x + w]
                try:
                    label, conf = self.face_recognizer.predict(face_roi)
                    if conf < self.conf_threshold:
                        user_id = self.label_map.get(label, "未知")
                        cv2.putText(frame, f"{user_id} ({100 - conf:.1f}%)",
                                    (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
                except:
                    pass

        # 转换图像格式
        h, w = frame.shape[:2]
        q_img = QImage(frame.data, w, h, 3 * w, QImage.Format_BGR888)
        self.video_label.setPixmap(QPixmap.fromImage(q_img).scaled(
            self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def toggle_play_pause(self):
        """优化后的播放/暂停逻辑"""
        self.is_paused = not self.is_paused
        if self.is_paused:
            self.play_pause_btn.setText("开始识别")
            self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-start"))
        else:
            self.play_pause_btn.setText("停止识别")
            self.play_pause_btn.setIcon(QIcon.fromTheme("media-playback-stop"))

    def update_threshold(self, value):
        """更新识别灵敏度"""
        self.conf_threshold = value
        self.status_label.setText(f"识别灵敏度：{self.conf_threshold}% | 状态：就绪")

    def load_dataset(self):
        """优化后的数据集加载"""
        images = []
        labels = []
        label_map = {}
        current_label = 0

        if not os.path.exists('pic'):
            os.makedirs('pic')

        for filename in os.listdir('pic'):
            if filename.endswith('.jpg'):
                id_num = filename.split('-')[0]
                if id_num not in label_map:
                    label_map[id_num] = current_label
                    current_label += 1
                label = label_map[id_num]

                img_path = os.path.join('pic', filename)
                img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                if img is not None:
                    img = cv2.resize(img, (100, 100))
                    images.append(img)
                    labels.append(label)

        if len(images) > 0:
            self.face_recognizer.train(images, np.array(labels))
        self.label_map = {v: k for k, v in label_map.items()}

    def capture_face(self):
        """优化后的人脸采集"""
        name = self.name_edit.text().strip()
        id_num = self.id_edit.text().strip()

        if not id_num:
            QMessageBox.warning(self, "输入错误", "学号不能为空！", QMessageBox.Ok)
            return

        ret, frame = self.cap.read()
        if not ret:
            QMessageBox.critical(self, "摄像头错误", "无法读取摄像头数据。", QMessageBox.Ok)
            return

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        if len(faces) == 0:
            QMessageBox.warning(self, "未检测到人脸", "请确保正对摄像头。", QMessageBox.Ok)
            return

        success_count = 0
        for i, (x, y, w, h) in enumerate(faces):
            face_img = gray[y:y + h, x:x + w]
            face_img = cv2.resize(face_img, (100, 100))
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
            file_name = f"pic/{id_num}-{timestamp}.jpg"
            cv2.imwrite(file_name, face_img)
            success_count += 1

        self.load_dataset()
        QMessageBox.information(self, "采集成功",
                                f"成功保存 {success_count} 张人脸图片！",
                                QMessageBox.Ok)

    def punch_in(self):
        """优化后的打卡逻辑"""
        if self.is_paused:
            QMessageBox.warning(self, "系统暂停", "请先开始识别。", QMessageBox.Ok)
            return

        ret, frame = self.cap.read()
        if not ret:
            QMessageBox.critical(self, "摄像头错误", "无法读取摄像头数据。", QMessageBox.Ok)
            return

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        if len(faces) == 0:
            QMessageBox.warning(self, "未检测到人脸", "请确保正对摄像头。", QMessageBox.Ok)
            return

        for (x, y, w, h) in faces:
            face_roi = gray[y:y + h, x:x + w]
            try:
                label, conf = self.face_recognizer.predict(face_roi)
                if conf < self.conf_threshold:
                    user_id = self.label_map.get(label, "未知")
                    if user_id == "未知":
                        raise ValueError("未注册用户")
                    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    self.save_punch_record(user_id, current_time)
                    QMessageBox.information(self, "打卡成功",
                                            f"{user_id} 于 {current_time} 打卡成功！",
                                            QMessageBox.Ok)
                    return
                else:
                    raise ValueError("识别失败")
            except:
                QMessageBox.warning(self, "识别失败", "未找到匹配的用户，请重新尝试。", QMessageBox.Ok)
                return

    def save_punch_record(self, user_id, time_str):
        """优化后的记录保存"""
        with open("attendance_records.csv", "a", encoding="utf-8-sig") as f:
            f.write(f"{user_id},{time_str}\n")

    def timerEvent(self, event):
        self.update_frame()

    def closeEvent(self, event):
        self.cap.release()
        event.accept()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    app.setStyle("Fusion")
    ex = FaceRecognitionSystem()
    ex.show()
    sys.exit(app.exec_())